# entrance of program
# author-by: xjtu-blacksmith
# create-on: 2020/6/14

import argparse as arg

from train import train_net
from model import vgg16_101

if __name__ == '__main__':
    parser = arg.ArgumentParser(description='Train a Pytorch CNN model.')
    parser.add_argument('--lr', metavar='e', nargs=1, type=int, help='the initial learning rate', default=0.01)
    parser.add_argument('--epoch', metavar='N', nargs=1, type=int, help='total epoch num', default=10)
    parser.add_argument('--label', metavar='text', nargs=1, help='label for this train', default='vgg16')

    args = parser.parse_args()
    net = vgg16_101(pretrained=True)
    train_net(net, lr=args.lr, epoch_num=args.epoch, msg=args.label)